from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import asyncio
from typing import Set, Dict, AsyncIterable
import logging

from service import VoiceBotService
from utils import *

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

app = FastAPI()

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该设置具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
active_connections: Set[WebSocket] = set()
voice_service: Dict[str, VoiceBotService] = {}

@app.get("/health")
async def health_check():
    return JSONResponse(content={"status": "ok"}, status_code=200)

@app.get("/actuator/health")
async def health_check():
    return JSONResponse(content={"status": "ok"}, status_code=200)

@app.websocket("/rtc/session/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    await websocket.accept()
    active_connections.add(websocket)
    
    try:
        # 创建并初始化VoiceBotService
        service = VoiceBotService(
            llm_ep_id="ep-20250306170957-9jbjt",
            tts_app_key="9460949811",
            tts_access_key="c8j4SlCFHFyqxhYGP5s_P4q8SEcilSKR",
            asr_app_key="9460949811",
            asr_access_key="c8j4SlCFHFyqxhYGP5s_P4q8SEcilSKR",
            session_id=session_id,
        )
        await service.init()
        voice_service[session_id] = service

        # 发送ready消息
        await websocket.send_bytes(
            convert_web_event_to_binary(
                WebEvent.from_payload(BotReadyPayload(session=session_id))
            )
        )

        async def async_gen(ws: WebSocket) -> AsyncIterable[WebEvent]:
            try:
                while True:
                    try:
                        message = await ws.receive_bytes()
                        input_event = convert_binary_to_web_event_to_binary(message)
                        # logging.info(
                        #     f"Received input event: {input_event.event}, "
                        #     f"payload: {input_event.event}, data len:{len(input_event.data)}"
                        # )
                        yield input_event
                    except WebSocketDisconnect:
                        logging.info("WebSocket disconnected")
                        break
                    except Exception as e:
                        logging.error(f"Error receiving message: {e}")
                        break
            except Exception as e:
                logging.error(f"Error in async_gen: {e}")

        async def fetch_output(ws: WebSocket, output_events: AsyncIterable[WebEvent]) -> None:
            try:
                async for output_event in output_events:
                    try:
                        logging.info(
                            f"Sending output event= {output_event.event}, "
                            f"data len:{len(output_event.data) if output_event.data else 0}, "
                            f"payload: {output_event.payload}"
                        )
                        await ws.send_bytes(convert_web_event_to_binary(output_event))
                    except WebSocketDisconnect:
                        logging.info("WebSocket disconnected while sending")
                        break
                    except Exception as e:
                        logging.error(f"Error sending message: {e}")
                        break
            except Exception as e:
                logging.error(f"Error in fetch_output: {e}")

        # 启动处理循环
        outputs = service.handler_loop(async_gen(websocket))
        await fetch_output(websocket, outputs)

    except WebSocketDisconnect:
        logging.info(f"WebSocket disconnected for session {session_id}")
    except Exception as e:
        logging.error(f"WebSocket error: {e}")
    finally:
        active_connections.remove(websocket)
        if session_id in voice_service:
            service = voice_service[session_id]
            del voice_service[session_id]
        try:
            await websocket.close()
        except:
            pass

async def start_http_server():
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=8080,
        ws_ping_interval=20,
        ws_ping_timeout=20,
        ws_max_size=10 * 1024 * 1024,  # 10MB
        log_level="info",
        access_log=True,
        proxy_headers=True,
        forwarded_allow_ips="*",
    )
    server = uvicorn.Server(config)
    await server.serve()

# 不再需要直接运行的部分
# if __name__ == "__main__":
#     start_server() 